import numpy as np
from backend.experiment.framework.lmmodel import LMModel


class PerplexityCalculator:
    def __init__(self, model: LMModel):
        self.__model = model

    @property
    def model(self) -> LMModel:
        return self.__model

    def perplexity(
            self, inputs: np.ndarray, begin: int = 0, end: int = None
    ) -> float:
        if len(inputs.shape) == 2:
            assert inputs.shape[0] == 1
            inputs = inputs.reshape((-1,))
        else:
            assert len(inputs.shape) == 1

        input_len = inputs.shape[0]
        assert input_len < self.__model.context_length

        if end is None:
            end = input_len

        assert end <= input_len

        p = []

        # 模型接收2阶形状的输入
        output_logits = self.__model.predict(inputs.reshape((1, -1)))
        if self.__model.output_logits:
            logits_exp = np.exp(output_logits)
            logits_exp_sum = np.sum(logits_exp, axis=-1)
            logits_exp_sum = np.expand_dims(
                logits_exp_sum, axis=logits_exp_sum.ndim)
            output_prob = np.squeeze(logits_exp / logits_exp_sum)
        else:
            output_prob = np.squeeze(output_logits)

        if len(output_prob.shape) == 1:
            output_prob = output_prob.reshape((1, -1))
        for i in range(begin, end - 1):
            # if i - 1 >= 0:
            #     print(f'inputs[i - 1]: {output_prob[i][inputs[i - 1]]}')
            # if i - 2 >= 0:
            #     print(f'inputs[i - 2]: {output_prob[i][inputs[i - 2]]}')
            # print(f'inputs[i]: {output_prob[i][inputs[i]]}')
            # if i + 1 < end - 1:
            #     print(f'inputs[i + 1]: {output_prob[i][inputs[i + 1]]}')
            # if i + 2 < end - 1:
            #     print(f'inputs[i + 2]: {output_prob[i][inputs[i + 2]]}')
            # print(f'mean: {np.mean(output_prob[i], axis=None)}')
            # print(f'max: {np.max(output_prob[i], axis=None)}')
            # print(f'min: {np.min(output_prob[i], axis=None)}')
            # input('::')
            p.append(np.log(output_prob[i][inputs[i + 1]] + 1e-6))
        # print(f'mean: {-1 * np.mean(p)}')
        # print(f'sum: {np.sum(p) / (begin - end - 1)}')
        return np.exp(-1 * np.mean(p))
